from conf.params_config import *
from business.crnn import CRNNHandle
from business.angnet import AngleNetHandle
from common.ocr_lite_utils import sorted_boxes, get_rotate_crop_image
from PIL import Image
import numpy as np
import copy
from business.dbnet_infer import DBNET
import traceback


class OcrHandle(object):
    def __init__(self):
        self.text_handle = DBNET(model_path)
        self.crnn_handle = CRNNHandle(crnn_model_path)
        if angle_detect:
            self.angle_handle = AngleNetHandle(angle_net_path)

    def crnnRecWithBox(self, im, boxes_list, score_list):
        """
        crnn模型，ocr识别
        @@model,
        @@converter,
        @@im:Array
        @@text_recs:text box
        @@ifIm:是否输出box对应的img

        """
        results = []
        boxes_list = sorted_boxes(np.array(boxes_list))

        line_imgs = []
        for index, (box, score) in enumerate(zip(boxes_list[:angle_detect_num], score_list[:angle_detect_num])):
            tmp_box = copy.deepcopy(box)
            partImg_array = get_rotate_crop_image(im, tmp_box.astype(np.float32))
            partImg = Image.fromarray(partImg_array).convert("RGB")
            line_imgs.append(partImg)

        angle_res = False
        if angle_detect:
            angle_res = self.angle_handle.predict_rbgs(line_imgs)
        if angle_res:
            return [], True
        count = 1
        for index, (box, score) in enumerate(zip(boxes_list, score_list)):
            tmp_box = copy.deepcopy(box)
            partImg_array = get_rotate_crop_image(im, tmp_box.astype(np.float32))

            partImg = Image.fromarray(partImg_array).convert("RGB")

            if angle_detect and angle_res:
                partImg = partImg.rotate(180)

            if not is_rgb:
                partImg = partImg.convert('L')

            try:

                if is_rgb:
                    simPred = self.crnn_handle.predict_rbg(partImg)  ##识别的文本
                else:
                    simPred = self.crnn_handle.predict(partImg)  ##识别的文本

            except Exception as e:
                print(traceback.format_exc())
                continue
            if simPred.strip() != '':
                tmp = dict()
                tmp["text"] = simPred.replace(" ", "").replace("　", "")
                bbox = tmp_box.tolist()
                tmp["score"] = round(float(score), 3)
                tmp["bbox"] = [
                    bbox[0][0], bbox[0][1], bbox[1][0], bbox[1][1],
                    bbox[2][0], bbox[2][1], bbox[3][0], bbox[3][1]
                ]
                results.append(tmp)
                count += 1

        return results, angle_res

    def text_predict(self, img, short_size):
        boxes_list, score_list = self.text_handle.process(np.asarray(img).astype(np.uint8), short_size=short_size)
        result, angle_res = self.crnnRecWithBox(np.array(img), boxes_list, score_list)
        return result, angle_res

    def text_rec(self, img):
        simPred = self.crnn_handle.predict_rbg(img)  ##识别的文本
        return simPred

